{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "xzpUtDMqmA-x"
      },
      "source": [
        "# Shakespeare in 5 minutes with Cloud TPUs\n",
        "\n",
        "This notebook demonstrates using Cloud TPUs to build a _language model_: a model that predicts the next character of text given the text so far.  Once our model has been trained we can sample from it to generate new text that looks like the text it was trained on.  In this case we're going to train our network using the combined works of Shakespeare, creating a play generating robot.\n",
        "\n",
        "\n",
        "### Note: You will need a GCP account and a GCS bucket for this notebook to run!\n",
        "\n",
        "Our network outputs something Shakespeare-esque:\n",
        "\n",
        "___\n",
        "\u003cblockquote\u003e\n",
        "Loves that led me no dumbs lack her Berjoy's face with her to-day.  \n",
        "The spirits roar'd; which shames which within his powers  \n",
        "\tWhich tied up remedies lending with occasion,  \n",
        "A loud and Lancaster, stabb'd in me  \n",
        "\tUpon my sword for ever: 'Agripo'er, his days let me free.  \n",
        "\tStop it of that word, be so: at Lear,  \n",
        "\tWhen I did profess the hour-stranger for my life,  \n",
        "\tWhen I did sink to be cried how for aught;  \n",
        "\tSome beds which seeks chaste senses prove burning;  \n",
        "But he perforces seen in her eyes so fast;  \n",
        "And _  \n",
        "\u003c/blockquote\u003e\n",
        "___\n",
        "\n",
        "Let's get started on generating our own Shakespeare!  We'll start off with our data generator.  The training data to our model will be snippets from our text file: the _target_ snippet is offset by one character.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "myGkRWgYWD2g"
      },
      "outputs": [],
      "source": [
        "# !rm /content/adc.json"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "both",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 259
        },
        "colab_type": "code",
        "id": "IcZkpa-e-Fas",
        "outputId": "f848b7fe-41a8-44a8-95bc-5297f0ac10ba"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Using bucket: tpu-estimator-shakespeare-test-bucket\n",
            "Using model dir: gs://tpu-estimator-shakespeare-test-bucket/tpuestimator-lstm/2018-09-28-23-58-37\n",
            "[_DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:CPU:0, CPU, -1, 15880407734472941098),\n",
            " _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 17179869184, 7578514533265224491),\n",
            " _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:XLA_GPU:0, XLA_GPU, 17179869184, 3512042959205926245),\n",
            " _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU:0, TPU, 17179869184, 6509007211901600635),\n",
            " _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU:1, TPU, 17179869184, 2788113998249947095),\n",
            " _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU:2, TPU, 17179869184, 18075511148356623033),\n",
            " _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU:3, TPU, 17179869184, 6450852309571070103),\n",
            " _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU:4, TPU, 17179869184, 14749604383048689166),\n",
            " _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU:5, TPU, 17179869184, 5384492138625106038),\n",
            " _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU:6, TPU, 17179869184, 17430458968359885752),\n",
            " _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU:7, TPU, 17179869184, 12326938441744536100),\n",
            " _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 17179869184, 2858742845697236993)]\n"
          ]
        }
      ],
      "source": [
        "import json\n",
        "import os\n",
        "import pprint\n",
        "import re\n",
        "import time\n",
        "import tensorflow as tf\n",
        "\n",
        "\n",
        "use_tpu = True #@param {type:\"boolean\"}\n",
        "bucket = '' #@param {type:\"string\"}\n",
        "\n",
        "assert bucket, 'Must specify an existing GCS bucket name'\n",
        "print('Using bucket: {}'.format(bucket))\n",
        "\n",
        "if use_tpu:\n",
        "    assert 'COLAB_TPU_ADDR' in os.environ, 'Missing TPU; did you request a TPU in Notebook Settings?'\n",
        "\n",
        "MODEL_DIR = 'gs://{}/{}'.format(bucket, time.strftime('tpuestimator-lstm/%Y-%m-%d-%H-%M-%S'))\n",
        "print('Using model dir: {}'.format(MODEL_DIR))\n",
        "\n",
        "from google.colab import auth\n",
        "auth.authenticate_user()\n",
        "\n",
        "if 'COLAB_TPU_ADDR' in os.environ:\n",
        "  TF_MASTER = 'grpc://{}'.format(os.environ['COLAB_TPU_ADDR'])\n",
        "  \n",
        "  # Upload credentials to TPU.\n",
        "  with tf.Session(TF_MASTER) as sess:    \n",
        "    with open('/content/adc.json', 'r') as f:\n",
        "      auth_info = json.load(f)\n",
        "    tf.contrib.cloud.configure_gcs(sess, credentials=auth_info)\n",
        "  # Now credentials are set for all future sessions on this TPU.\n",
        "else:\n",
        "  TF_MASTER=''\n",
        "\n",
        "with tf.Session(TF_MASTER) as session:\n",
        "  pprint.pprint(session.list_devices())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Qew6Qt0-kGzO"
      },
      "source": [
        "# Training Data\n",
        "\n",
        "We can use a `tf.data` pipeline to feed input data to our Estimator.  In this case, we want our model to predict the next character, so we will feed sequences from our dataset where the source is offset from the target by 1 character.\n",
        "\n",
        "Note that we use `tf.contrib.data.enumerate_dataset()` and  `tf.contrib.stateless.stateless_random_uniform` to generate deterministic uniform samples.  This, combined with the setting of `RunConfig.tf_random_seed` guarantees that every run of the model will have the exact same behavior."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 173
        },
        "colab_type": "code",
        "id": "E3V4V-Jxmuv3",
        "outputId": "0ea9525a-764c-4d33-fd5c-1bf7efabcc15"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Redirecting output to ‘wget-log.1’.\n",
            "Batch size: 1\n",
            "INFO:tensorflow:Sample text: ureless, and rude, barrenly perish:\n",
            "Look whom she best endowed, she gave thee more;\n",
            "Which bounteou\n",
            "Tensor(\"Cast:0\", shape=(), dtype=int32)\n",
            "[[111 109 105 115 101 100  39 115 116  32]]\n",
            "[[109 105 115 101 100  39 115 116  32 116]]\n"
          ]
        }
      ],
      "source": [
        "import numpy as np\n",
        "\n",
        "!wget --show-progress --continue -O /content/shakespeare.txt http://www.gutenberg.org/files/100/100-0.txt\n",
        "\n",
        "SHAKESPEARE_TXT = '/content/shakespeare.txt'\n",
        "RANDOM_SEED = 42  # An arbitrary choice.\n",
        "\n",
        "def transform(txt):\n",
        "  return np.asarray([ord(c) for c in txt], dtype=np.int32)\n",
        "\n",
        "def input_fn(params):\n",
        "  \"\"\"Return a dataset of source and target sequences for training.\"\"\"\n",
        "  batch_size = params['batch_size']\n",
        "  print('Batch size: {}'.format(batch_size))\n",
        "  seq_len = params['seq_len']\n",
        "  with tf.gfile.GFile(params['source_file'], 'r') as f:\n",
        "    txt = f.read()\n",
        "    txt = ''.join([x for x in txt if ord(x) \u003c 128])\n",
        "    \n",
        "  tf.logging.info('Sample text: %s', txt[10000:10100])\n",
        "  source = tf.constant(transform(txt), dtype=tf.int32)\n",
        "  ds = tf.data.Dataset.from_tensors(source)\n",
        "  ds = ds.repeat()\n",
        "  ds = ds.apply(tf.contrib.data.enumerate_dataset())\n",
        "\n",
        "  def _select_seq(offset, src):\n",
        "    idx = tf.contrib.stateless.stateless_random_uniform(\n",
        "        [1], seed=[RANDOM_SEED, offset], dtype=tf.float32)[0]\n",
        "\n",
        "    max_start_offset = len(txt) - seq_len\n",
        "    idx = tf.cast(idx * max_start_offset, tf.int32)\n",
        "    print(idx)\n",
        "    \n",
        "    return {\n",
        "        'source': tf.reshape(src[idx:idx + seq_len], [seq_len]),\n",
        "        'target': tf.reshape(src[idx + 1:idx + seq_len + 1], [seq_len])\n",
        "    }\n",
        "\n",
        "  ds = ds.map(_select_seq)\n",
        "  ds = ds.batch(batch_size, drop_remainder=True)\n",
        "  ds = ds.prefetch(2)\n",
        "  return ds\n",
        "\n",
        "tf.reset_default_graph()\n",
        "tf.set_random_seed(0)\n",
        "with tf.Session() as session:\n",
        "  ds = input_fn({'batch_size': 1, 'seq_len': 10, 'source_file': SHAKESPEARE_TXT})\n",
        "  features = session.run(ds.make_one_shot_iterator().get_next())\n",
        "  print(features['source'])\n",
        "  print(features['target'])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Bbb05dNynDrQ"
      },
      "source": [
        "## Building our model\n",
        "\n",
        "Now that we have some data, we can define our model.  We use a simple 3 layer, forward LSTM for this notebook.\n",
        "\n",
        "The only change to our model versus a CPU/GPU model is that we specify a static `shape` for the input of our model.  This allows TF to infer the shape of the model and satisfy the XLA compilers static shape requirement."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "yLEM-fLJlEEt"
      },
      "outputs": [],
      "source": [
        "EMBEDDING_DIM = 1024\n",
        "\n",
        "# Construct a 2-layer LSTM\n",
        "def _lstm(inputs, batch_size, initial_state=None):\n",
        "  def _make_cell(layer_idx):\n",
        "    with tf.variable_scope('lstm/%d' % layer_idx,):\n",
        "      return tf.nn.rnn_cell.LSTMCell(\n",
        "          num_units=EMBEDDING_DIM,\n",
        "          state_is_tuple=True,\n",
        "          reuse=tf.AUTO_REUSE,\n",
        "      )\n",
        "\n",
        "  cell = tf.nn.rnn_cell.MultiRNNCell([\n",
        "      _make_cell(0), \n",
        "      _make_cell(1),\n",
        "  ])\n",
        "  if initial_state is None:\n",
        "    initial_state = cell.zero_state(batch_size, tf.float32)\n",
        "\n",
        "  outputs, final_state = tf.contrib.recurrent.functional_rnn(\n",
        "      cell, inputs, initial_state=initial_state, use_tpu=use_tpu)\n",
        "  return outputs, final_state\n",
        "\n",
        "\n",
        "def lstm_model(seq, initial_state=None):\n",
        "  with tf.variable_scope('lstm', \n",
        "                         initializer=tf.orthogonal_initializer,\n",
        "                         reuse=tf.AUTO_REUSE):\n",
        "    batch_size = seq.shape[0]\n",
        "    seq_len = seq.shape[1]\n",
        "\n",
        "    embedding_params = tf.get_variable(\n",
        "        'char_embedding', \n",
        "        initializer=tf.orthogonal_initializer(seed=0),\n",
        "        shape=(256, EMBEDDING_DIM), dtype=tf.float32)\n",
        "\n",
        "    embedding = tf.nn.embedding_lookup(embedding_params, seq)\n",
        "\n",
        "    lstm_output, lstm_state = _lstm(\n",
        "        embedding, batch_size, initial_state=initial_state)\n",
        "\n",
        "    # Apply a single dense layer to the output of our LSTM to predict\n",
        "    # our final characters.  This looks awkward as we have to flatten\n",
        "    # our input to 2 dimensions before applying the dense layer.\n",
        "    flattened = tf.reshape(lstm_output, [-1, EMBEDDING_DIM])\n",
        "    logits = tf.layers.dense(flattened, 256, name='logits',)\n",
        "    logits = tf.reshape(logits, [-1, seq_len, 256])\n",
        "    return logits, lstm_state"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "j0ZYOd07qJws"
      },
      "source": [
        "## Training our model\n",
        "\n",
        "Since we're using TPUEstimator, we need to provide what's called a _model function_ to train our model.  This specifies how to train, evaluate and run inference (predictions) on our model.\n",
        "\n",
        "Let's cover each part in turn.  We'll first look at the training step.  \n",
        "\n",
        "* We feed our source tensor to our LSTM model\n",
        "* Compute the cross entropy loss to train it better predict the target tensor.\n",
        "* Use the `RMSPropOptimizer` to optimize our network\n",
        "* Wrap it with the `CrossShardOptimizer` which lets us use multiple TPU cores to train.  \n",
        "\n",
        "Finally we return a `TPUEstimatorSpec` indicating how TPUEstimator should train our model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "1b5E8ZSUrCBk"
      },
      "outputs": [],
      "source": [
        "def train_fn(source, target):\n",
        "  logits, lstm_state = lstm_model(source)\n",
        "  batch_size = source.shape[0]\n",
        "  \n",
        "  loss = tf.reduce_mean(\n",
        "      tf.nn.sparse_softmax_cross_entropy_with_logits(\n",
        "          labels=target, logits=logits))\n",
        "\n",
        "  optimizer = tf.train.AdamOptimizer(learning_rate=0.001)\n",
        "  if TF_MASTER:\n",
        "    optimizer = tf.contrib.tpu.CrossShardOptimizer(optimizer)\n",
        "  train_op = optimizer.minimize(loss, tf.train.get_global_step())\n",
        "  return tf.contrib.tpu.TPUEstimatorSpec(\n",
        "      mode=tf.estimator.ModeKeys.TRAIN,\n",
        "      loss=loss,\n",
        "      train_op=train_op,\n",
        "  )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ro-Y2oG27l4r"
      },
      "source": [
        "## Evaluating our model\n",
        "\n",
        "Next, evaluation.  This is simpler: we run our model forward and check how well it predicts the next character.  Again, we return a `TPUEstimatorSpec` to tell TPUEstimator how to evaluate the model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "gycj1IPp63Fj"
      },
      "outputs": [],
      "source": [
        "def eval_fn(source, target):\n",
        "  logits, _ = lstm_model(source)\n",
        "  def metric_fn(labels, logits):\n",
        "    labels = tf.cast(labels, tf.int64)\n",
        "    return {\n",
        "        'recall@1': tf.metrics.recall_at_k(labels, logits, 1),\n",
        "        'recall@5': tf.metrics.recall_at_k(labels, logits, 5)\n",
        "    }\n",
        "\n",
        "  eval_metrics = (metric_fn, [target, logits])\n",
        "  return tf.contrib.tpu.TPUEstimatorSpec(\n",
        "      mode=tf.estimator.ModeKeys.EVAL, \n",
        "      loss=loss, \n",
        "      eval_metrics=eval_metrics)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "rHenJneu78Sy"
      },
      "source": [
        "## Computing Predictions\n",
        "\n",
        "We leave the most complicated part for last.  There's nothing TPU specific here!  For predictions we use the input tensor as a _seed_ for our model.  We then use a TensorFlow loop to sample characters from our model and return the result."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "-bRhg5Tx8PLr"
      },
      "outputs": [],
      "source": [
        "def predict_fn(source):\n",
        "  # Seed the model with our initial array\n",
        "  batch_size = source.shape[0]\n",
        "  logits, lstm_state = lstm_model(source)\n",
        "\n",
        "  def _body(i, state, preds):\n",
        "    \"\"\"Body of our prediction loop: predict the next character.\"\"\"\n",
        "    cur_preds = preds.read(i)\n",
        "    next_logits, next_state = lstm_model(\n",
        "        tf.cast(tf.expand_dims(cur_preds, -1), tf.int32), state)\n",
        "\n",
        "    # pull out the last (and only) prediction.\n",
        "    next_logits = next_logits[:, -1]\n",
        "    next_pred = tf.multinomial(\n",
        "        next_logits, num_samples=1, output_dtype=tf.int32)[:, 0]\n",
        "    preds = preds.write(i + 1, next_pred)\n",
        "    return (i + 1, next_state, preds)\n",
        "\n",
        "  def _cond(i, state, preds):\n",
        "    del state\n",
        "    del preds\n",
        "\n",
        "    # Loop until `predict_len - 1`: preds[0] is the initial state and we\n",
        "    # write to `i + 1` on each iteration.\n",
        "    return tf.less(i, predict_len - 1)\n",
        "\n",
        "  next_pred = tf.multinomial(\n",
        "      logits[:, -1], num_samples=1, output_dtype=tf.int32)[:, 0]\n",
        "\n",
        "  i = tf.constant(0, dtype=tf.int32)\n",
        "\n",
        "  predict_len = 500\n",
        "\n",
        "  # compute predictions as [seq_len, batch_size] to simplify indexing/updates\n",
        "  pred_var = tf.TensorArray(\n",
        "      dtype=tf.int32,\n",
        "      size=predict_len,\n",
        "      dynamic_size=False,\n",
        "      clear_after_read=False,\n",
        "      element_shape=(batch_size,),\n",
        "      name='prediction_accumulator',\n",
        "  )\n",
        "\n",
        "  pred_var = pred_var.write(0, next_pred)\n",
        "  _, _, final_predictions = tf.while_loop(_cond, _body,\n",
        "                                          [i, lstm_state, pred_var])\n",
        "\n",
        "  # reshape back to [batch_size, predict_len] and cast to int32\n",
        "  final_predictions = final_predictions.stack()\n",
        "  final_predictions = tf.transpose(final_predictions, [1, 0])\n",
        "  final_predictions = tf.reshape(final_predictions, (batch_size, predict_len))\n",
        "\n",
        "  return tf.contrib.tpu.TPUEstimatorSpec(\n",
        "      mode=tf.estimator.ModeKeys.PREDICT, \n",
        "      predictions={'predictions': final_predictions})"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "IwRTkF4l8e3M"
      },
      "source": [
        "## Building our model function\n",
        "\n",
        "We can now use our helper functions to build our combined model function and train our model!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "D5e8TD4q8kq2"
      },
      "outputs": [],
      "source": [
        "def model_fn(features, labels, mode, params):\n",
        "  if mode == tf.estimator.ModeKeys.TRAIN:\n",
        "    return train_fn(features['source'], features['target'])\n",
        "  if mode == tf.estimator.ModeKeys.EVAL:\n",
        "    return eval_fn(features['source'], features['target'])\n",
        "  if mode == tf.estimator.ModeKeys.PREDICT:\n",
        "    return predict_fn(features['source'])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "VzBYDJI0_Tfm"
      },
      "source": [
        "## Running our model\n",
        "\n",
        "We now have a bit of boilerplate to specify our TPU worker and then we can train our model!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 2369
        },
        "colab_type": "code",
        "id": "CS9no3m_rCf0",
        "outputId": "87bfd24a-ef1d-4a65-9b2d-0d18438eceda"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Using config: {'_model_dir': 'gs://tpu-estimator-shakespeare-test-bucket/tpuestimator-lstm/2018-09-28-23-58-37', '_tf_random_seed': 42, '_save_summary_steps': 100, '_save_checkpoints_steps': 5000, '_save_checkpoints_secs': None, '_session_config': allow_soft_placement: true\n",
            "graph_options {\n",
            "  rewrite_options {\n",
            "    meta_optimizer_iterations: ONE\n",
            "  }\n",
            "}\n",
            ", '_keep_checkpoint_max': 5, '_keep_checkpoint_every_n_hours': 10000, '_log_step_count_steps': None, '_train_distribute': None, '_device_fn': None, '_protocol': None, '_eval_distribute': None, '_experimental_distribute': None, '_service': None, '_cluster_spec': \u003ctensorflow.python.training.server_lib.ClusterSpec object at 0x7f2d308f1ef0\u003e, '_task_type': 'worker', '_task_id': 0, '_global_id_in_cluster': 0, '_master': 'grpc://10.76.7.218:8470', '_evaluation_master': 'grpc://10.76.7.218:8470', '_is_chief': True, '_num_ps_replicas': 0, '_num_worker_replicas': 1, '_tpu_config': TPUConfig(iterations_per_loop=100, num_shards=8, num_cores_per_replica=None, per_host_input_for_training=2, tpu_job_name=None, initial_infeed_sleep_secs=None, input_partition_dims=None), '_cluster': None}\n",
            "INFO:tensorflow:_TPUContext: eval_on_tpu True\n",
            "INFO:tensorflow:Querying Tensorflow master (grpc://10.76.7.218:8470) for TPU system metadata.\n",
            "INFO:tensorflow:Found TPU system:\n",
            "INFO:tensorflow:*** Num TPU Cores: 8\n",
            "INFO:tensorflow:*** Num TPU Workers: 1\n",
            "INFO:tensorflow:*** Num TPU Cores Per Worker: 8\n",
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:CPU:0, CPU, -1, 15880407734472941098)\n",
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 17179869184, 7578514533265224491)\n",
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:XLA_GPU:0, XLA_GPU, 17179869184, 3512042959205926245)\n",
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU:0, TPU, 17179869184, 6509007211901600635)\n",
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU:1, TPU, 17179869184, 2788113998249947095)\n",
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU:2, TPU, 17179869184, 18075511148356623033)\n",
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU:3, TPU, 17179869184, 6450852309571070103)\n",
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU:4, TPU, 17179869184, 14749604383048689166)\n",
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU:5, TPU, 17179869184, 5384492138625106038)\n",
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU:6, TPU, 17179869184, 17430458968359885752)\n",
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU:7, TPU, 17179869184, 12326938441744536100)\n",
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 17179869184, 2858742845697236993)\n",
            "INFO:tensorflow:Calling model_fn.\n",
            "Batch size: 1024\n",
            "INFO:tensorflow:Sample text: ureless, and rude, barrenly perish:\n",
            "Look whom she best endowed, she gave thee more;\n",
            "Which bounteou\n",
            "Tensor(\"Cast:0\", shape=(), dtype=int32, device=/job:tpu_worker/task:0/device:CPU:0)\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/gradients_impl.py:108: UserWarning: Converting sparse IndexedSlices to a dense Tensor of unknown shape. This may consume a large amount of memory.\n",
            "  \"Converting sparse IndexedSlices to a dense Tensor of unknown shape. \"\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Create CheckpointSaverHook.\n",
            "INFO:tensorflow:Done calling model_fn.\n",
            "INFO:tensorflow:TPU job name tpu_worker\n",
            "INFO:tensorflow:Graph was finalized.\n",
            "INFO:tensorflow:Restoring parameters from gs://tpu-estimator-shakespeare-test-bucket/tpuestimator-lstm/2018-09-28-23-58-37/model.ckpt-500\n",
            "INFO:tensorflow:Running local_init_op.\n",
            "INFO:tensorflow:Done running local_init_op.\n",
            "INFO:tensorflow:Saving checkpoints for 500 into gs://tpu-estimator-shakespeare-test-bucket/tpuestimator-lstm/2018-09-28-23-58-37/model.ckpt.\n",
            "INFO:tensorflow:Installing graceful shutdown hook.\n",
            "INFO:tensorflow:Creating heartbeat manager for ['/job:tpu_worker/replica:0/task:0/device:CPU:0', '/job:tpu_worker/replica:0/task:0/device:XLA_CPU:0']\n",
            "WARNING:tensorflow:Worker heartbeats not supported by all workers.  No failure handling will be enabled.\n",
            "INFO:tensorflow:Init TPU system\n",
            "INFO:tensorflow:Starting infeed thread controller.\n",
            "INFO:tensorflow:Starting outfeed thread controller.\n",
            "INFO:tensorflow:Enqueue next (100) batch(es) of data to infeed.\n",
            "INFO:tensorflow:Dequeue next (100) batch(es) of data from outfeed.\n",
            "INFO:tensorflow:loss = 1.2207099, step = 600\n",
            "INFO:tensorflow:Enqueue next (100) batch(es) of data to infeed.\n",
            "INFO:tensorflow:Dequeue next (100) batch(es) of data from outfeed.\n",
            "INFO:tensorflow:loss = 1.2004857, step = 700 (16.893 sec)\n",
            "INFO:tensorflow:global_step/sec: 5.91939\n",
            "INFO:tensorflow:examples/sec: 6061.46\n",
            "INFO:tensorflow:Enqueue next (100) batch(es) of data to infeed.\n",
            "INFO:tensorflow:Dequeue next (100) batch(es) of data from outfeed.\n",
            "INFO:tensorflow:loss = 1.1948798, step = 800 (16.768 sec)\n",
            "INFO:tensorflow:global_step/sec: 5.96381\n",
            "INFO:tensorflow:examples/sec: 6106.95\n",
            "INFO:tensorflow:Enqueue next (100) batch(es) of data to infeed.\n",
            "INFO:tensorflow:Dequeue next (100) batch(es) of data from outfeed.\n",
            "INFO:tensorflow:loss = 1.1526046, step = 900 (19.553 sec)\n",
            "INFO:tensorflow:global_step/sec: 5.11454\n",
            "INFO:tensorflow:examples/sec: 5237.29\n",
            "INFO:tensorflow:Enqueue next (100) batch(es) of data to infeed.\n",
            "INFO:tensorflow:Dequeue next (100) batch(es) of data from outfeed.\n",
            "INFO:tensorflow:loss = 1.1549901, step = 1000 (16.756 sec)\n",
            "INFO:tensorflow:global_step/sec: 5.96787\n",
            "INFO:tensorflow:examples/sec: 6111.1\n",
            "INFO:tensorflow:Enqueue next (100) batch(es) of data to infeed.\n",
            "INFO:tensorflow:Dequeue next (100) batch(es) of data from outfeed.\n",
            "INFO:tensorflow:loss = 1.1079394, step = 1100 (16.774 sec)\n",
            "INFO:tensorflow:global_step/sec: 5.96144\n",
            "INFO:tensorflow:examples/sec: 6104.52\n",
            "INFO:tensorflow:Enqueue next (100) batch(es) of data to infeed.\n",
            "INFO:tensorflow:Dequeue next (100) batch(es) of data from outfeed.\n",
            "INFO:tensorflow:loss = 1.123244, step = 1200 (16.742 sec)\n",
            "INFO:tensorflow:global_step/sec: 5.97283\n",
            "INFO:tensorflow:examples/sec: 6116.18\n",
            "INFO:tensorflow:Enqueue next (100) batch(es) of data to infeed.\n",
            "INFO:tensorflow:Dequeue next (100) batch(es) of data from outfeed.\n",
            "INFO:tensorflow:loss = 1.0875183, step = 1300 (16.782 sec)\n",
            "INFO:tensorflow:global_step/sec: 5.95897\n",
            "INFO:tensorflow:examples/sec: 6101.98\n",
            "INFO:tensorflow:Enqueue next (100) batch(es) of data to infeed.\n",
            "INFO:tensorflow:Dequeue next (100) batch(es) of data from outfeed.\n",
            "INFO:tensorflow:loss = 1.0983433, step = 1400 (16.738 sec)\n",
            "INFO:tensorflow:global_step/sec: 5.97466\n",
            "INFO:tensorflow:examples/sec: 6118.05\n",
            "INFO:tensorflow:Enqueue next (100) batch(es) of data to infeed.\n",
            "INFO:tensorflow:Dequeue next (100) batch(es) of data from outfeed.\n",
            "INFO:tensorflow:loss = 1.0734489, step = 1500 (16.777 sec)\n",
            "INFO:tensorflow:global_step/sec: 5.96048\n",
            "INFO:tensorflow:examples/sec: 6103.53\n",
            "INFO:tensorflow:Enqueue next (100) batch(es) of data to infeed.\n",
            "INFO:tensorflow:Dequeue next (100) batch(es) of data from outfeed.\n",
            "INFO:tensorflow:loss = 1.049539, step = 1600 (16.748 sec)\n",
            "INFO:tensorflow:global_step/sec: 5.97095\n",
            "INFO:tensorflow:examples/sec: 6114.25\n",
            "INFO:tensorflow:Enqueue next (100) batch(es) of data to infeed.\n",
            "INFO:tensorflow:Dequeue next (100) batch(es) of data from outfeed.\n",
            "INFO:tensorflow:loss = 1.026352, step = 1700 (19.171 sec)\n",
            "INFO:tensorflow:global_step/sec: 5.21613\n",
            "INFO:tensorflow:examples/sec: 5341.32\n",
            "INFO:tensorflow:Enqueue next (100) batch(es) of data to infeed.\n",
            "INFO:tensorflow:Dequeue next (100) batch(es) of data from outfeed.\n",
            "INFO:tensorflow:loss = 1.0030547, step = 1800 (16.763 sec)\n",
            "INFO:tensorflow:global_step/sec: 5.96552\n",
            "INFO:tensorflow:examples/sec: 6108.69\n",
            "INFO:tensorflow:Enqueue next (100) batch(es) of data to infeed.\n",
            "INFO:tensorflow:Dequeue next (100) batch(es) of data from outfeed.\n",
            "INFO:tensorflow:loss = 1.0061597, step = 1900 (16.776 sec)\n",
            "INFO:tensorflow:global_step/sec: 5.96088\n",
            "INFO:tensorflow:examples/sec: 6103.94\n",
            "INFO:tensorflow:Enqueue next (100) batch(es) of data to infeed.\n",
            "INFO:tensorflow:Dequeue next (100) batch(es) of data from outfeed.\n",
            "INFO:tensorflow:loss = 0.9461493, step = 2000 (16.768 sec)\n",
            "INFO:tensorflow:global_step/sec: 5.96372\n",
            "INFO:tensorflow:examples/sec: 6106.85\n",
            "INFO:tensorflow:Saving checkpoints for 2000 into gs://tpu-estimator-shakespeare-test-bucket/tpuestimator-lstm/2018-09-28-23-58-37/model.ckpt.\n",
            "INFO:tensorflow:Stop infeed thread controller\n",
            "INFO:tensorflow:Shutting down InfeedController thread.\n",
            "INFO:tensorflow:InfeedController received shutdown signal, stopping.\n",
            "INFO:tensorflow:Infeed thread finished, shutting down.\n",
            "INFO:tensorflow:infeed marked as finished\n",
            "INFO:tensorflow:Stop output thread controller\n",
            "INFO:tensorflow:Shutting down OutfeedController thread.\n",
            "INFO:tensorflow:OutfeedController received shutdown signal, stopping.\n",
            "INFO:tensorflow:Outfeed thread finished, shutting down.\n",
            "INFO:tensorflow:outfeed marked as finished\n",
            "INFO:tensorflow:Shutdown TPU system.\n",
            "INFO:tensorflow:Loss for final step: 0.9461493.\n",
            "INFO:tensorflow:training_loop marked as finished\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "\u003ctensorflow.contrib.tpu.python.tpu.tpu_estimator.TPUEstimator at 0x7f2d307910f0\u003e"
            ]
          },
          "execution_count": 26,
          "metadata": {
            "tags": []
          },
          "output_type": "execute_result"
        }
      ],
      "source": [
        "def _make_estimator(num_shards, use_tpu=True):\n",
        "  config = tf.contrib.tpu.RunConfig(\n",
        "      tf_random_seed=RANDOM_SEED,\n",
        "      master=TF_MASTER,\n",
        "      model_dir=MODEL_DIR,\n",
        "      save_checkpoints_steps=5000,\n",
        "      tpu_config=tf.contrib.tpu.TPUConfig(\n",
        "          num_shards=num_shards, iterations_per_loop=100))\n",
        "\n",
        "  estimator = tf.contrib.tpu.TPUEstimator(\n",
        "      use_tpu=use_tpu,\n",
        "      model_fn=model_fn, config=config,\n",
        "      train_batch_size=1024,\n",
        "      eval_batch_size=1024,\n",
        "      predict_batch_size=128,\n",
        "      params={'seq_len': 100, 'source_file': SHAKESPEARE_TXT},\n",
        "  )\n",
        "  return estimator\n",
        "\n",
        "\n",
        "# Use all 8 cores for training\n",
        "estimator = _make_estimator(num_shards=8, use_tpu=use_tpu)\n",
        "estimator.train(\n",
        "    input_fn=input_fn,\n",
        "    max_steps=2000,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "TCBtcpZkykSf"
      },
      "source": [
        "## Running predictions with our model\n",
        "\n",
        "We've trained our model, now we can run predictions through it to generate \"Shakespeare\"!  We provide a seed sentence to get our model started, and then sample 500 characters from it."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 555
        },
        "colab_type": "code",
        "id": "tU7M-EGGxR3E",
        "outputId": "d57d39bb-c0c6-4f3a-ab78-c7c9258bd2a2"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Using config: {'_model_dir': 'gs://tpu-estimator-shakespeare-test-bucket/tpuestimator-lstm/2018-09-28-23-58-37', '_tf_random_seed': 42, '_save_summary_steps': 100, '_save_checkpoints_steps': 5000, '_save_checkpoints_secs': None, '_session_config': allow_soft_placement: true\n",
            "graph_options {\n",
            "  rewrite_options {\n",
            "    meta_optimizer_iterations: ONE\n",
            "  }\n",
            "}\n",
            ", '_keep_checkpoint_max': 5, '_keep_checkpoint_every_n_hours': 10000, '_log_step_count_steps': None, '_train_distribute': None, '_device_fn': None, '_protocol': None, '_eval_distribute': None, '_experimental_distribute': None, '_service': None, '_cluster_spec': \u003ctensorflow.python.training.server_lib.ClusterSpec object at 0x7f2d30618f60\u003e, '_task_type': 'worker', '_task_id': 0, '_global_id_in_cluster': 0, '_master': 'grpc://10.76.7.218:8470', '_evaluation_master': 'grpc://10.76.7.218:8470', '_is_chief': True, '_num_ps_replicas': 0, '_num_worker_replicas': 1, '_tpu_config': TPUConfig(iterations_per_loop=100, num_shards=1, num_cores_per_replica=None, per_host_input_for_training=2, tpu_job_name=None, initial_infeed_sleep_secs=None, input_partition_dims=None), '_cluster': None}\n",
            "WARNING:tensorflow:Setting TPUConfig.num_shards==1 is an unsupported behavior. Please fix as soon as possible (leaving num_shards as None.\n",
            "INFO:tensorflow:_TPUContext: eval_on_tpu True\n",
            "WARNING:tensorflow:eval_on_tpu ignored because use_tpu is False.\n",
            "INFO:tensorflow:Calling model_fn.\n",
            "INFO:tensorflow:Running infer on CPU\n",
            "INFO:tensorflow:Done calling model_fn.\n",
            "INFO:tensorflow:Graph was finalized.\n",
            "INFO:tensorflow:Restoring parameters from gs://tpu-estimator-shakespeare-test-bucket/tpuestimator-lstm/2018-09-28-23-58-37/model.ckpt-2000\n",
            "INFO:tensorflow:Running local_init_op.\n",
            "INFO:tensorflow:Done running local_init_op.\n",
            "INFO:tensorflow:prediction_loop marked as finished\n",
            "One forward I am.\n",
            "  DEMETRIUS. Sir, O, be sin to me and more defence,\n",
            "    Made wars shall pinch this large request again\n",
            "    Whiles my young mistress Mercury.\n",
            "  ROSALIND. 'Tis a coward that carries them.\n",
            "    By the garter, he shall roast thy book, must pleasure, kneel'd;\n",
            "    you'll be dient,\n",
            "    with my son Lucius.\n",
            "  CORIOLANUS. Hear me; if thou read be won order than too.\n",
            "    The long thou seest is lost; but out thy taints\n",
            "    Look forth thee well, if more, a perisher, go.\n",
            "    Chirti\n"
          ]
        }
      ],
      "source": [
        "def _seed_input_fn(params):\n",
        "  del params\n",
        "  seed_txt = 'Looks it not like the king?'\n",
        "  seed = transform(seed_txt)\n",
        "  seed = tf.constant(seed.reshape([1, -1]), dtype=tf.int32)\n",
        "  # Predict must return a Dataset, not a Tensor.\n",
        "  return tf.data.Dataset.from_tensors({'source': seed})\n",
        "\n",
        "# Use 1 core for prediction since we're only generating a single element batch\n",
        "estimator = _make_estimator(num_shards=1, use_tpu=False)\n",
        "\n",
        "idx = next(estimator.predict(input_fn=_seed_input_fn))['predictions']\n",
        "print(''.join([chr(i) for i in idx]))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "4vCt_FSTeBQq"
      },
      "outputs": [],
      "source": [
        ""
      ]
    }
  ],
  "metadata": {
    "accelerator": "TPU",
    "colab": {
      "collapsed_sections": [],
      "name": "Predict Shakespeare with Cloud TPUs and TPUEstimator",
      "provenance": [],
      "version": "0.3.2"
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
